//
// Created by lancelot on 17-7-13.
//

#include "Optimizer.h"
#include "../../DataStructure/Frame.h"
#include "../../DataStructure/LandMarkPLX.h"
#include "DataManager.h"
#include <ceres/ceres.h>
#include "../include/PBAFactors.h"
// #include "ceres/CeresPBAFactor4Nodes.h"
#include "../include/PBABlocks.hpp"
#include <cstring>
#include <fstream>
#include <iostream>

namespace Estimator
{
using namespace Eigen;

void Optimizer::AddPoseBlock(const Frame* pframe)
{
    Eigen::Matrix<double, 3, 4, RowMajor> frameTransMat;
    frameTransMat.topLeftCorner<3,3>() = pframe->mRotationMatrix;
    frameTransMat.topRightCorner<3,1>() = pframe->mTranslationVec;

    // save it into Pose Vector
    // mvPoseBlock.emplace_back(pframe->GetFrameId(), frameTransMat);
    mHashFrames.emplace(std::piecewise_construct, std::forward_as_tuple(pframe->GetFrameId()),
                                                      std::forward_as_tuple(pframe->GetFrameId(), frameTransMat));
    mOptimizationProblem.AddParameterBlock(mHashFrames.at(pframe->GetFrameId()).parameters(), PoseBlock::Dimension);
    mOptimizationProblem.SetParameterization( mHashFrames[pframe->GetFrameId()].parameters(), mPoseLocalPtr_);
}

void Optimizer::AddLandMarkBlock(const Estimator::LandMarkPLX* landPLX)
{
    //std::cout << "== enter add landmark ====" << std::endl;
   
    Eigen::Matrix<double, 5, 1 > featureParallax;
    featureParallax(0) = landPLX->GetCosParallax();;
    featureParallax(1) = landPLX->GetSinParallax();
    featureParallax.segment<3>(2) = landPLX->GetDirectionMain();

    // save it into LandMark Vector

    mHashLandMarks.emplace(std::piecewise_construct, std::forward_as_tuple(landPLX->GetId()), 
                                                    std::forward_as_tuple(landPLX->GetId(), featureParallax));
    mOptimizationProblem.AddParameterBlock( mHashLandMarks[landPLX->GetId()].parameters(), LandMarkBlock::Dimension);
    mOptimizationProblem.SetParameterization( mHashLandMarks[landPLX->GetId()].parameters(), mpLandMarkLocalPtr_);
}

void Optimizer::GenerateCostFunctionBeta(const Managing* pManaging)
{
    for(const auto landMark : pManaging->mlLandMarkContainer_)
    {
        const Frame *pMainFrame = landMark.GetMainFrame();
        const Frame *pAssoFrame = landMark.GetAssoFrame();

        if(pAssoFrame == nullptr)
            continue;

        auto mapObservations = landMark.GetObservations();

        for(auto iteObservations = mapObservations.cbegin(); iteObservations!= mapObservations.cend(); iteObservations++)
        {
            if(iteObservations->first == pMainFrame) continue;
            else if(iteObservations->first == pAssoFrame)
            {
                Eigen::Matrix3d information = Eigen::Matrix3d::Identity();
                // Add three Nodes Edge
                auto correspondence = iteObservations->first->mvLandMarkObservations_[iteObservations->second];
                // std::array<double,2> uv = correspondence.second;
                Eigen::Vector3d temp = mCalibration.Kinv* Eigen::Vector3d(correspondence.second[0], correspondence.second[1] ,1);
                temp.normalize();

                ceres::CostFunction* cost_function = new CeresPBAFactor3Nodes(temp, information);
                mOptimizationProblem.AddResidualBlock(cost_function, NULL, 
                                                        mHashFrames[pMainFrame->GetFrameId()].parameters(), 
                                                        mHashFrames[pAssoFrame->GetFrameId()].parameters(),
                                                        mHashLandMarks[landMark.GetId()].parameters());
            }
            else
            {
                // Add four nodes edge.
                Eigen::Matrix3d information=Eigen::Matrix3d::Identity();
                // Add three Nodes Edge
                auto correspondence = iteObservations->first->mvLandMarkObservations_[iteObservations->second];

                Eigen::Vector3d temp = mCalibration.Kinv* Eigen::Vector3d(correspondence.second[0], correspondence.second[1] ,1);
                temp.normalize();

                ceres::CostFunction* cost_function = new CeresPBAFactor4Nodes(temp, information);
                mOptimizationProblem.AddResidualBlock(cost_function, NULL, 
                                                        mHashFrames[pMainFrame->GetFrameId()].parameters(), 
                                                        mHashFrames[pAssoFrame->GetFrameId()].parameters(), 
                                                        mHashFrames[iteObservations->first->GetFrameId()].parameters(),
                                                        mHashLandMarks[landMark.GetId()].parameters());
            }
        }
    }


    mOptimizationProblem.SetParameterBlockConstant(mHashFrames[0].parameters());
    mOptimizationProblem.SetParameterBlockConstant(mHashFrames[1].parameters());
    SetOrdering();
}

void Optimizer::SetOrdering()
{
    ceres::ParameterBlockOrdering *ordering = new ceres::ParameterBlockOrdering;
    // The points come before the cameras

    for(auto ite = mHashLandMarks.begin(); ite!= mHashLandMarks.end(); ite++){
        ordering->AddElementToGroup((ite->second).parameters(), 0);
    }

    for(auto ite = mHashFrames.begin(); ite!= mHashFrames.end(); ite++)
    {
        ordering->AddElementToGroup((ite->second).parameters(), 1);
    }
    mOptions_.linear_solver_ordering.reset(ordering);
}

void Optimizer::SolveOptimization() 
{
    ceres::Solver::Summary summary;
    ceres::Solve(mOptions_, &mOptimizationProblem, &summary);
    std::cout << summary.FullReport() << "\n";
}

void Optimizer::InitializeByManager(const Managing* pManager)
{
    mHashFrames.reserve(pManager->mvpFrameContainer_.size());
    mHashLandMarks.reserve(pManager->mlLandMarkContainer_.size());
    // AddPoseBlock(pFrame);
    for(auto pframe: pManager->mvpFrameContainer_)
    {
        AddPoseBlock(pframe);
    }

    for(auto LandMark : pManager->mlLandMarkContainer_)
    {
        AddLandMarkBlock(&LandMark);
    }
}

// void Optimizer::WriteResult()
// {}

// void Optimizer::WriteResultintoFile(std::string resultfile)
// {
//     this->WriteResult();
//     std::fstream file(resultfile + "poseValue.txt", std::fstream::out | std::fstream::trunc);
//     for (size_t i = 0; i < mvPoseGroundTruth.size(); ++i)
//     {
//         file << mvPoseGroundTruth[i].first << " ";
//         for (size_t j = 0; j < 12; ++j)
//         {
//             file << mvPoseGroundTruth[i].second.at(j) << " ";
//         }
//         file << std::endl;
//     }

//     std::fstream landmarkfile(resultfile + "LandMarkValue.txt", std::fstream::out | std::fstream::trunc);
//     for (size_t i = 0; i < mvLandMarkGroundTruth.size(); ++i)
//     {
//         landmarkfile << mvLandMarkGroundTruth[i].first << " ";
//         for(size_t j = 0; j < 5; ++j)
//         {
//             landmarkfile << mvLandMarkGroundTruth[i].second.at(j) << " ";
//         }
//         landmarkfile << std::endl;
//     }

//     file.close();
//     landmarkfile.close();
// }

void Optimizer::SetOptions()
{
    mOptions_.minimizer_type = ceres::TRUST_REGION;
    mOptions_.linear_solver_type = ceres::DENSE_SCHUR;
    // mOptions_.linear_solver_type = ceres::SPARSE_SCHUR;
    // mOptions_.linear_solver_type = ceres::ITERATIVE_SCHUR;
    // mOptions_.dense_linear_algebra_library_type = ceres::LAPACK; 


    mOptions_.trust_region_strategy_type = ceres::LEVENBERG_MARQUARDT;
    // mOptions_.minimizer_progress_to_stdout = true;
     
    mOptions_.max_num_iterations = 100;
    // mOptions_.dogleg_type = ceres::SUBSPACE_DOGLEG;
    // mOptions_.use_nonmonotonic_steps = true;
    //mOptions_.sparse_linear_algebra_library_type = ceres::SUITE_SPARSE;
    mOptions_.num_threads = 4;
    mOptions_.num_linear_solver_threads = 4;

    // mOptions_.callbacks.push_back(new StateUpdateCallback(this,"../Iteration", 1));
}
}